#include <Wire.h>
#include <Arduino.h>
#include <U8g2lib.h>

U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, /* reset=*/U8X8_PIN_NONE, /* SCL=*/22, /* SDA=*/21);

// 定义滤波增强参数
#define USE_SIGNAL_ENHANCEMENT true  // 是否启用信号增强
const int SMOOTH_WINDOW = 3;         // 平滑窗口大小
const float ALPHA = 0.2;             // 低通滤波系数

// HRV分析相关配置
const int RR_BUFFER_SIZE = 400;      // 5分钟数据存储
unsigned int rrIntervals[RR_BUFFER_SIZE];
int rrIndex = 0;
bool bufferFull = false;
unsigned long lastRPeakTime = 0;

// 心率有效性检查参数 
const int HR_MIN = 40;               // 最小合理心率
const int HR_MAX = 180;              // 最大合理心率
const int HR_HISTORY_SIZE = 5;       // 历史心率记录数
byte hrHistory[HR_HISTORY_SIZE];     // 心率历史记录
int hrHistoryIndex = 0;              // 心率历史记录索引
byte lastValidHR = 70;               // 上次有效心率

// Pan-Tompkins算法参数
#define ECG_BUFFER_SIZE 64      // ECG缓冲区大小，增大以适应更高采样率
#define SAMPLE_FREQ 512         // 采样频率512Hz
const float LOW_PASS_ALPHA = 0.08;     // 低通滤波系数，调整以适应更高采样率
const float HIGH_PASS_ALPHA = 0.85;    // 高通滤波系数，调整以适应更高采样率
const int INTEGRATION_WIDTH = 40;     // 积分窗口宽度，增大以适应更高采样率
const float DETECTION_THRESHOLD = 0.7; // 检测阈值系数

// Pan-Tompkins算法缓冲区
short ecgBuffer[ECG_BUFFER_SIZE];       // 原始ECG数据缓冲区
float lpfBuffer[ECG_BUFFER_SIZE];       // 低通滤波缓冲区
float hpfBuffer[ECG_BUFFER_SIZE];       // 高通滤波缓冲区
float derivBuffer[ECG_BUFFER_SIZE];     // 微分缓冲区
float squaredBuffer[ECG_BUFFER_SIZE];   // 平方缓冲区
float integralBuffer[ECG_BUFFER_SIZE];  // 积分缓冲区

int bufferIndex = 0;                    // 当前缓冲区索引
float threshold1 = 0;                   // 信号阈值
float threshold2 = 0;                   // 噪声阈值
float spki = 0;                         // 信号峰值积分
float npki = 0;                         // 噪声峰值积分
unsigned long lastRRInterval = 600;     // 上一次RR间隔(ms)
int rWaveCount = 0;                     // R波计数

// BMD101通信配置
#define BAUDRATE 57600
byte errorRate = 0;
byte heartrate = 0;
short raw;
boolean newRawData = false;
boolean bigPacket = false;

// 显示刷新控制
unsigned long lastDisplayUpdate = 0;

// 波形追踪和调试变量
unsigned long lastDataTime = 0;
boolean dataFlowing = false;

// 心率检测可靠性增强
unsigned long lastTwoRPeaks[2] = {0, 0}; // 存储最近两次R波时间
const unsigned long MIN_RR_GAP = 300;    // 最小合理RR间隔(ms)
const unsigned long MAX_MISSED_BEATS = 3; // 最大连续漏拍次数

void OLED_Logo() {
    u8g2.clearBuffer();
    u8g2.setFont(u8g2_font_ncenB08_tr);
    u8g2.drawStr(0, 10, "BMD101 Demo");
    u8g2.sendBuffer();
}

void OLED_InitDisplay() {
    u8g2.clearBuffer();
    u8g2.setFont(u8g2_font_ncenB08_tr);
    u8g2.drawStr(0, 10, "Initializing...");
    u8g2.sendBuffer();
}

// 心率平滑处理和有效性检查
byte validateHeartRate(byte newHR) {
    // 先检查是否是明显不合理的值
    if (newHR < HR_MIN || newHR > HR_MAX) {
        return lastValidHR; // 返回上次的有效心率
    }
    
    // 存储心率历史
    hrHistory[hrHistoryIndex] = newHR;
    hrHistoryIndex = (hrHistoryIndex + 1) % HR_HISTORY_SIZE;
    
    // 计算平均心率 (移动平均)
    int sum = 0;
    for (int i = 0; i < HR_HISTORY_SIZE; i++) {
        sum += hrHistory[i];
    }
    byte avgHR = sum / HR_HISTORY_SIZE;
    
    // 检查当前心率是否与平均值偏差太大
    if (abs(newHR - avgHR) > avgHR * 0.3) {
        // 偏差超过30%，可能是异常值
        return lastValidHR;
    }
    
    // 更新上次有效心率
    lastValidHR = newHR;
    return newHR;
}

void setup() {
    Serial.begin(BAUDRATE);
    u8g2.begin();
    u8g2.enableUTF8Print();
    OLED_Logo();
    delay(3000);
    OLED_InitDisplay();
    
    // 初始化Pan-Tompkins算法缓冲区
    for (int i = 0; i < ECG_BUFFER_SIZE; i++) {
        ecgBuffer[i] = 0;
        lpfBuffer[i] = 0;
        hpfBuffer[i] = 0;
        derivBuffer[i] = 0;
        squaredBuffer[i] = 0;
        integralBuffer[i] = 0;
    }
    
    // 初始化心率历史记录
    for (int i = 0; i < HR_HISTORY_SIZE; i++) {
        hrHistory[i] = 70; // 初始化为合理默认值
    }
    
    lastDataTime = millis();
}

void loop() {
    ReadData();
    
    // 检查数据流状态
    if (newRawData) {
        // 对原始信号进行增强处理
        short processedSample = USE_SIGNAL_ENHANCEMENT ? enhanceSignal(raw) : raw;
        
        // 处理信号
        processECGSample(processedSample);
        newRawData = false;
        dataFlowing = true;
        lastDataTime = millis();
    } else if (millis() - lastDataTime > 3000) {
        // 如果3秒内没有新数据，认为数据流中断
        dataFlowing = false;
    }
    
    // 检测漏拍情况，超过2秒没有检测到心跳
    if (lastRPeakTime > 0 && millis() - lastRPeakTime > 2000) {
        // 可能漏拍，使用上次有效心率
        heartrate = lastValidHR;
    }
    
    updateDisplay();
}

// Pan-Tompkins算法实现
void processECGSample(short ecgSample) {
    // 存储ECG样本
    ecgBuffer[bufferIndex] = ecgSample;
    
    // 低通滤波 - 实现公式: y(n) = α * x(n) + (1-α) * y(n-1)
    int prevIndex = (bufferIndex > 0) ? (bufferIndex - 1) : (ECG_BUFFER_SIZE - 1);
    lpfBuffer[bufferIndex] = LOW_PASS_ALPHA * ecgSample + (1 - LOW_PASS_ALPHA) * lpfBuffer[prevIndex];
    
    // 高通滤波 - 实现公式: y(n) = α * y(n-1) + α * (x(n) - x(n-1))
    hpfBuffer[bufferIndex] = HIGH_PASS_ALPHA * hpfBuffer[prevIndex] + 
                             HIGH_PASS_ALPHA * (lpfBuffer[bufferIndex] - lpfBuffer[prevIndex]);
    
    // 微分 - 五点差分方程: y(n) = (1/8) * [2x(n) + x(n-1) - x(n-3) - 2x(n-4)]
    int idx1 = (bufferIndex > 0) ? (bufferIndex - 1) : (ECG_BUFFER_SIZE - 1);
    int idx2 = (bufferIndex > 1) ? (bufferIndex - 2) : (ECG_BUFFER_SIZE - 2 + bufferIndex);
    int idx3 = (bufferIndex > 2) ? (bufferIndex - 3) : (ECG_BUFFER_SIZE - 3 + bufferIndex);
    int idx4 = (bufferIndex > 3) ? (bufferIndex - 4) : (ECG_BUFFER_SIZE - 4 + bufferIndex);
    
    derivBuffer[bufferIndex] = (2 * hpfBuffer[bufferIndex] + hpfBuffer[idx1] - 
                               hpfBuffer[idx3] - 2 * hpfBuffer[idx4]) / 8.0;
    
    // 平方
    squaredBuffer[bufferIndex] = derivBuffer[bufferIndex] * derivBuffer[bufferIndex];
    
    // 移动窗口积分
    float sum = 0;
    for (int i = 0; i < INTEGRATION_WIDTH; i++) {
        int idx = (bufferIndex - i + ECG_BUFFER_SIZE) % ECG_BUFFER_SIZE;
        sum += squaredBuffer[idx];
    }
    integralBuffer[bufferIndex] = sum / INTEGRATION_WIDTH;
    
    // 自适应阈值QRS检测
    detectRWave(integralBuffer[bufferIndex]);
    
    // 更新缓冲区索引
    bufferIndex = (bufferIndex + 1) % ECG_BUFFER_SIZE;
}

// R波检测和RR间隔测量
void detectRWave(float integralValue) {
    // 初始化阈值
    if (rWaveCount == 0) {
        threshold1 = integralValue / 2;
        threshold2 = threshold1 / 2;
        spki = threshold1;
        npki = threshold2;
    }
    
    // 搜索窗口内的峰值检测
    static boolean lookForRWave = true;
    static unsigned long lastPeakTime = 0;
    static float peakValue = 0;
    static unsigned long peakTime = 0;
    unsigned long currentTime = millis();
    
    // 确保足够时间过去以避免T波检测（不少于160ms - 适应512Hz采样率）
    if ((currentTime - lastPeakTime) < 160) {
        return;
    }
    
    // 峰值检测
    if (integralValue > peakValue && lookForRWave) {
        peakValue = integralValue;
        peakTime = currentTime;
    }
    
    // 确定该峰值是否为R波（扫描窗口调整为40ms，适应512Hz）
    if (lookForRWave && integralValue < peakValue && (currentTime - peakTime) > 40) {
        lookForRWave = false;
        
        // 根据阈值判断是否为R波
        if (peakValue > threshold1) {
            // 是R波
            unsigned long RRInterval = peakTime - lastRPeakTime;
            
            // 更新最近两次R波时间
            lastTwoRPeaks[0] = lastTwoRPeaks[1];
            lastTwoRPeaks[1] = peakTime;
            
            // 检测是否可能漏检了心跳
            if (RRInterval > 1200 && lastRRInterval > 0 && lastRRInterval < 1000) {
                // 如果当前RR间隔异常长且上次RR间隔正常，可能漏检了心跳
                // 估计漏检的心跳次数
                int missedBeats = round(RRInterval / lastRRInterval) - 1;
                
                if (missedBeats > 0 && missedBeats <= MAX_MISSED_BEATS) {
                    // 估计漏检的心跳，维持之前的心率
                    heartrate = lastValidHR;
                    
                    // 不要将这个异常长的RR间隔加入HRV计算
                    RRInterval = 0; // 标记为无效
                }
            }
            
            // 始终更新心率，无论RR间隔是否在范围内
            if (RRInterval > MIN_RR_GAP) {
                byte calculatedHR = 60000 / RRInterval;
                // 使用心率验证函数处理
                heartrate = validateHeartRate(calculatedHR);
            }
            
            // 存储有效RR间隔，使用更严格的有效性检查
            if (RRInterval > MIN_RR_GAP && RRInterval < 1500) {
                // 检查最近的两次RR间隔，如果变化太大，可能是异常
                if (lastTwoRPeaks[0] > 0) {
                    unsigned long prevRR = lastTwoRPeaks[1] - lastTwoRPeaks[0];
                    float ratio = (float)RRInterval / prevRR;
                    
                    // 如果相邻RR间隔变化超过40%，认为是不稳定的
                    if (ratio >= 0.6 && ratio <= 1.4) {
                        // 存储有效RR间隔
                        rrIntervals[rrIndex] = RRInterval;
                        rrIndex = (rrIndex + 1) % RR_BUFFER_SIZE;
                        if (!bufferFull && rrIndex == 0) {
                            bufferFull = true;
                        }
                        
                        // 更新上一次R波时间和RR间隔
                        lastRRInterval = RRInterval;
                    }
                }
                else {
                    // 第一次检测，直接存储
                    rrIntervals[rrIndex] = RRInterval;
                    rrIndex = (rrIndex + 1) % RR_BUFFER_SIZE;
                    if (!bufferFull && rrIndex == 0) {
                        bufferFull = true;
                    }
                    
                    // 更新上一次R波时间和RR间隔
                    lastRRInterval = RRInterval;
                }
            }
            
            // 始终更新上一次R波时间，即使RR间隔不在范围内
            lastRPeakTime = peakTime;
            
            // 更新信号阈值 - 动态调整以适应不同幅度的信号
            float spkiFactor = errorRate > 100 ? 0.15 : 0.125; // 高噪声时更快适应
            spki = spkiFactor * peakValue + (1 - spkiFactor) * spki;
            threshold1 = npki + 0.25 * (spki - npki);
            threshold2 = threshold1 / 2;
            
            rWaveCount++;
        } else if (peakValue > threshold2) {
            // 是噪声峰值
            npki = 0.125 * peakValue + 0.875 * npki;
            threshold1 = npki + 0.25 * (spki - npki);
            threshold2 = threshold1 / 2;
        }
        
        // 重置峰值搜索
        peakValue = 0;
        lastPeakTime = currentTime;
    }
    
    // 重置R波搜索标志
    if (!lookForRWave && integralValue < threshold2) {
        lookForRWave = true;
    }
}

// 心率变异分析函数
float calculateSDNN() {
    int count = bufferFull ? RR_BUFFER_SIZE : rrIndex;
    if (count < 2) return 0;

    float sum = 0.0, mean, sdnn = 0.0;
    for (int i = 0; i < count; ++i) sum += rrIntervals[i];
    mean = sum / count;

    for (int i = 0; i < count; ++i)
        sdnn += pow(rrIntervals[i] - mean, 2);

    return sqrt(sdnn / (count - 1));
}

// 计算均方根差
float calculateRMSSD() {
    int count = bufferFull ? RR_BUFFER_SIZE : rrIndex;
    if (count < 2) return 0;

    float sum = 0.0;
    for (int i = 1; i < count; ++i) {
        long diff = rrIntervals[i] - rrIntervals[i - 1];
        sum += pow(diff, 2);
    }
    return sqrt(sum / (count - 1));
}

// 计算pNN50
float calculatePNN50() {
    int count = bufferFull ? RR_BUFFER_SIZE : rrIndex;
    if (count < 2) return 0;

    int nn50 = 0;
    for (int i = 1; i < count; ++i) {
        unsigned int diff = abs((int)(rrIntervals[i] - rrIntervals[i - 1]));
        if (diff > 50) nn50++;
    }
    return (nn50 * 100.0) / (count - 1);
}

// 显示界面更新
void updateDisplay() {
    if (millis() - lastDisplayUpdate < 1000) return;

    u8g2.clearBuffer();

    // 第一行显示基础信息和数据流状态
    u8g2.setFont(u8g2_font_profont12_tf);
    u8g2.setCursor(0, 12);
    u8g2.print("S:");
    u8g2.print(errorRate);
    u8g2.setCursor(64, 12);
    u8g2.print("HR:");
    u8g2.print(heartrate);
    
    // 显示数据流状态
    u8g2.setCursor(110, 12);
    u8g2.print(dataFlowing ? "*" : "X");
    
    // 显示收集的数据点数
    u8g2.setCursor(0, 24);
    u8g2.print("RR:");
    u8g2.print(bufferFull ? RR_BUFFER_SIZE : rrIndex);
    
    // 显示Pan-Tompkins算法状态
    u8g2.setCursor(64, 24);
    u8g2.print("R:");
    u8g2.print(rWaveCount);

    // 显示HRV分析结果，降低需要的样本数量为50
    if (rrIndex > 50 || bufferFull) {
        u8g2.setCursor(0, 34);
        u8g2.print("SDNN:");
        u8g2.print(calculateSDNN(), 1);

        u8g2.setCursor(0, 44);
        u8g2.print("RMSSD:");
        u8g2.print(calculateRMSSD(), 1);

        u8g2.setCursor(0, 54);
        u8g2.print("pNN50:");
        u8g2.print(calculatePNN50(), 1);
        u8g2.print("%");
    } else {
        u8g2.setFont(u8g2_font_ncenB10_tr);
        u8g2.drawStr(20, 40, "Collecting...");
    }

    // 绘制进度条 - 确保可见
    int progress = map(rrIndex, 0, RR_BUFFER_SIZE, 0, 100);
    u8g2.drawFrame(0, 58, 128, 6);
    int barWidth = (int)(progress * 1.28);
    if (barWidth < 1) barWidth = 1;
    u8g2.drawBox(0, 58, barWidth, 6); // 确保至少有1像素宽度

    u8g2.sendBuffer();
    lastDisplayUpdate = millis();
}

// BMD101数据读取函数
byte ReadOneByte() {
    while (Serial.available() < 1);
    return Serial.read();
}

void ReadData() {
    static unsigned char payloadData[256]; // 用于存储数据包有效载荷的缓冲区
    byte generatedChecksum; // 计算得到的校验和
    byte checksum; // 从数据包中读取的校验和
    byte vLength; // 数据长度
    int payloadLength; // 有效载荷长度

    // 检查是否收到两个连续的0xAA（十进制170）字节，BMD101数据包的起始标记
    if (ReadOneByte() == 170) {
        if (ReadOneByte() == 170) {
            // 读取数据包的长度。如果读到的值仍然是170，则继续读取下一个字节，直到获取到非170的值。如果长度超过169，则认为数据包无效并退出函数。
            do {
                payloadLength = ReadOneByte();
            } while (payloadLength == 170);
            if (payloadLength > 169) {
                return;
            }

            // 读取有效载荷数据并计算校验和。BMD101使用的校验和算法是将所有有效载荷字节相加，然后取255减去这个和
            generatedChecksum = 0;
            for (int i = 0; i < payloadLength; i++) {
                payloadData[i] = ReadOneByte(); // 读取有效载荷数据
                generatedChecksum += payloadData[i]; // 累加计算校验和
            }
            checksum = ReadOneByte(); // 读取数据包中的校验和
            generatedChecksum = 255 - generatedChecksum; // 计算校验和的补码

            // 校验和不匹配，表示数据包可能在传输过程中损坏。
            if (checksum != generatedChecksum) {
                // checksum error
                // 这里进行校验失败操作
            } else {
                // 数据包解析
                for (int i = 0; i < payloadLength; i++) {
                    switch (payloadData[i]) {
                        case 2: // 信号质量数据
                            bigPacket = true;
                            i++;
                            errorRate = payloadData[i]; // 存储信号质量值
                            break;
                        case 3: // 心率数据
                            i++;
                            // 只在无Pan-Tompkins算法时使用BMD101心率
                            // 如果信号质量为0，不更新心率，避免出现心率为0
                            if (rWaveCount == 0 && errorRate > 0) {
                                heartrate = payloadData[i]; 
                            }
                            i = i + 14; // 跳过14个字节
                            break;
                        case 0x80: // 原始ECG数据
                            newRawData = true;
                            i++;
                            vLength = payloadData[i]; // 数据长度
                            raw = 0;
                            for (int j = 0; j < vLength; j++) {
                                raw = raw | (payloadData[i + vLength - j] << (8 * j)); // 组合多个字节为一个16位的值
                            }
                            i += vLength;
                            break;
                        default:
                            break;
                    }
                }
            }
        }
    }
}

// 信号增强滤波函数
short enhanceSignal(short rawSignal) {
    static short smoothBuffer[SMOOTH_WINDOW];
    static float lastFiltered = 0;
    static int smoothIndex = 0;
    
    // 简单的移动平均滤波，减少高频噪声
    smoothBuffer[smoothIndex] = rawSignal;
    smoothIndex = (smoothIndex + 1) % SMOOTH_WINDOW;
    
    int sum = 0;
    for (int i = 0; i < SMOOTH_WINDOW; i++) {
        sum += smoothBuffer[i];
    }
    short smoothed = sum / SMOOTH_WINDOW;
    
    // 一阶低通滤波，进一步平滑信号
    float filtered = ALPHA * smoothed + (1 - ALPHA) * lastFiltered;
    lastFiltered = filtered;
    
    return (short)filtered;
}